import json
from typing import Dict, List
import mysql.connector
from mysql.connector import Error

class KnowledgeGraph:
    def __init__(self, db_config: Dict):
        self.db_config = db_config
        self.connection = self._create_connection()

    def _create_connection(self):
        """创建数据库连接"""
        try:
            return mysql.connector.connect(
                host=self.db_config['host'],
                port=self.db_config['port'],
                user=self.db_config['user'],
                password=self.db_config['password'],
                database=self.db_config['database']
            )
        except Error as e:
            print(f"Error while connecting to MySQL: {e}")
            return None

    def get_pet_knowledge(self, pet_type: str) -> List[Dict]:
        """根据宠物类型获取知识图谱数据"""
        query = f"SELECT * FROM {self.db_config['dialogue_table']} WHERE pet_type = %s"
        cursor = self.connection.cursor(dictionary=True)
        cursor.execute(query, (pet_type,))
        return cursor.fetchall()

    def save_knowledge(self, knowledge: Dict):
        """保存知识图谱数据"""
        query = f"INSERT INTO {self.db_config['dialogue_table']} (pet_type, question, answer) VALUES (%s, %s, %s)"
        cursor = self.connection.cursor()
        cursor.execute(query, (knowledge['pet_type'], knowledge['question'], knowledge['answer']))
        self.connection.commit()

    def close(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()

if __name__ == "__main__":
    # 示例用法
    config = {
        'host': 'rm-2zea4lxfy2974xc9xso.mysql.rds.aliyuncs.com',
        'port': 3306,
        'user': 'ry',
        'password': 'Qwer1111',
        'database': 'mingxiangzhu_mall',
        'dialogue_table': 'pet_ai_talk'
    }
    
    kg = KnowledgeGraph(config)
    try:
        # 获取狗狗相关知识
        dog_knowledge = kg.get_pet_knowledge('dog')
        print(json.dumps(dog_knowledge, indent=2, ensure_ascii=False))
        
        # 保存新的知识
        new_knowledge = {
            'pet_type': 'cat',
            'question': '猫需要每天梳毛吗？',
            'answer': '是的，特别是长毛猫需要每天梳毛以防止毛发打结。'
        }
        kg.save_knowledge(new_knowledge)
    finally:
        kg.close()
